import copy
import time

from get_sim_embeddings_m3e import  GetEmbedding
import numpy as np
import json
import faiss
import random


#获取embeddings
model_path = './finetuned-model/model/'
# model_path = '/data/huggingface/models--moka-ai--m3e-base/snapshots/764b537a0e50e5c7d64db883f2d2e051cbe3c64c/'
# # test set
# embedding_path = "./dataset/test_set_v10.1.npy"
# data_path = './dataset/test_set_v10.1.json'
# all set
embedding_path = "./dataset/train_test_set_v10.1.npy"
data_path = './dataset/train_test_set_v10.1.json'

# data_path = './dataset/train_test_set_v1.json'
print(f'model_path: {model_path}')
print(f'embedding_path: {embedding_path}')
# 建立查询矩阵
test = GetEmbedding(model_path)
embeddings_ = np.load(embedding_path)
di_qu_all_embedding = np.array(embeddings_)
quantizer = faiss.IndexFlatIP(di_qu_all_embedding.shape[1])
quantizer.add(di_qu_all_embedding.astype(np.float32))

def readjson(file_path):
  """"""
  # make_path_legal(file_path)
  with open(file_path, "r", encoding="utf-8") as f:
    return json.load(f)

# datas = readjson('./symptoms_all.json') all_answer_men_te_man
datas = readjson(data_path)
# datas = readjson('./dataset/all_answer_men_te_man.json')
data_keys = []
data_vals = []
for k,v, prefix in datas:
   data_keys.append(k)
#    data_vals.append(f'{prefix}@{v}')
   data_vals.append(v)

def get_augument_content(input_text, topk=5):
    print("问题embeddding时间：")
    start_e = time.time()
    query_embeding = test.get_sim_embedding(input_text)
    end_e = time.time()
    print("问题embeddding时间：",str(end_e - start_e))

    #query_embeding_ori = copy.copy(query_embeding)
    # normalize(query_embeding)
    print("-----------------------------------")    
    print(query_embeding.shape)
    start_r = time.time()
    distance, idx = quantizer.search(query_embeding, topk)
    end_r = time.time()
    print("问题检索时间：",str(end_r - start_r))
    print("distance: ", distance)
    print("idx: ", idx)
    #答案list
    q_annd_a = {}
    for index,item in enumerate(idx[0]):
        #print(index,item,di_qu_all_question[item],"\t",distance[0][index])#di_qu_name[item])
        ask_ = data_keys[item]
        que_ = data_vals[item]
            #answer_list.append(da_an)
        if len(ask_) > 1:
            q_annd_a.update({
                str(index):[ask_, que_, str(distance[0][index])]
            })

    return  q_annd_a

def valid_test_set():
    top1_cts = 0
    top5_cts = 0

    corr_dis = []
    error_dis = []

    for q, a, *_ in datas:
        sear = get_augument_content(q)
        print(q, a, sear)
        top5 = False
        top1 = False
        dis = float(sear['0'][2])
        for i in range(5):
            q_q, q_a, _ = sear[str(i)]
            if q_a == a:
                top5 = True
                if i==0:
                    top1 = True
                    print(f'correct:' + '*' * 20)
        if top5:
            top5_cts += 1
        if top1:
            top1_cts += 1
            corr_dis.append(dis)
        else:
            error_dis.append(dis)

        # print(f'疾病检索：')
        # disease = q.split(' ')[0]
        # dis_sae = get_augument_content(disease)
        # print(disease, a, dis_sae)
        print('\n' + '=' * 80 + '\n')
    print(f'top1: {top1_cts/len(datas):.4f}')
    print(f'top5: {top5_cts/len(datas):.4f}')

    corr_dis = sorted(corr_dis)
    error_dis = sorted(error_dis)
    print(f'corr_dis: {corr_dis}')
    print(f'error_dis: {error_dis}')


def valid_all_set():
    top1_cts = 0
    top5_cts = 0

    corr_dis = []
    error_dis = []

    for q, a, *_ in datas:
        sear = get_augument_content(q)
        print(q, a, sear)
        top5 = False
        top1 = False
        dis = float(sear['0'][2])
        for i in range(5):
            q_q, q_a, _ = sear[str(i)]
            if q_a == a:
                top5 = True
                if i==0:
                    top1 = True
                    print(f'correct:' + '*' * 20)
        if top5:
            top5_cts += 1
        if top1:
            top1_cts += 1
            corr_dis.append(dis)
        else:
            error_dis.append(dis)

        # print(f'疾病检索：')
        # disease = q.split(' ')[0]
        # dis_sae = get_augument_content(disease)
        # print(disease, a, dis_sae)
        print('\n' + '=' * 80 + '\n')
    print(f'top1: {top1_cts/len(datas):.4f}')
    print(f'top5: {top5_cts/len(datas):.4f}')

    corr_dis = sorted(corr_dis)
    error_dis = sorted(error_dis)
    print(f'corr_dis: {corr_dis}')
    print(f'error_dis: {error_dis}')

def test_model():
#     test_datas = ['恶性肿瘤康复治疗 做过哪些治疗？']
#     test_datas = ['恶性肿瘤康复治疗 您尝试过哪些治疗方法？',
#  '恶性肿瘤康复治疗 您接受过哪些治疗？',
#  '恶性肿瘤康复治疗 您已经尝试过哪些治疗了？',
#  '恶性肿瘤康复治疗 您曾经尝试过哪些治疗方式？',
#  '恶性肿瘤康复治疗 您做过哪些治疗呢？',
#  '恶性肿瘤康复治疗 您已经尝试了哪些治疗方法？',
#  '恶性肿瘤康复治疗 您进行过哪些治疗？',
#  '恶性肿瘤康复治疗 您尝试了哪些治疗？',
#  '恶性肿瘤康复治疗 您曾经进行过哪些治疗？',
#  '恶性肿瘤康复治疗 您已经尝试过哪些治疗措施？',
#  '恶性肿瘤康复治疗 您做过哪些治疗方案？',
#  '恶性肿瘤康复治疗 您已经接受过哪些治疗？',
#  '恶性肿瘤康复治疗 您已经尝试了哪些治疗方案？',
#  '恶性肿瘤康复治疗 您进行了哪些治疗？',
#  '恶性肿瘤康复治疗 您已经尝试了哪些治疗方式？',
#  '恶性肿瘤康复治疗 您曾经接受过哪些治疗方式？',
#  '恶性肿瘤康复治疗 您接受了哪些治疗？',
#  '恶性肿瘤康复治疗 您已经进行过哪些治疗了？',
#  '恶性肿瘤康复治疗 您尝试过哪些治疗措施？',
#  '恶性肿瘤康复治疗 您已经进行了哪些治疗方法？',
#  '恶性肿瘤康复治疗 我感冒了']
    # test_datas = ['恶性肿瘤康复治疗 您好，请问您发现恶性肿瘤多久了？']
    test_datas = ['高血压3级 收缩压（SBP，血测压计左血压得边的高压）最高多少？',
                  '高血压3级 舒张压（血测压计右血压得边的低压）最高多少？']
    for data in test_datas:
        print(f'{data}： {get_augument_content(data)}')
    # for _ in range(100):
    #     i = random.randint(0, len(datas))
    #     data = datas[i]
    #     cur = get_augument_content(data)
    #     print(data, cur)

if __name__ == '__main__':
    # valid_test_set()

    valid_all_set()

    # test()
